from ..utils.database import InferSessionLocal
from ..model.infer import Material, Lot, AOI, Failure, InferPanel, InferPanelImage
from sqlalchemy import and_, or_
import json
import os
from datetime import datetime, timedelta
from ..config import Config

root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def get_material_name(material_code_str):
    with InferSessionLocal() as session:
        material_codes = material_code_str.split(", ")
        material_names = []
        for material_code in material_codes:
            material = session.query(Material).filter_by(material_code=material_code).first()
            if material:
                material_name = material.material_name
                material_names.append(material_name)
        material_name_str = material_names[0] if len(material_names)==1 else ", ".join(material_names)

        return material_name_str
    

def get_aoi_name(aoi_code_str):
    with InferSessionLocal() as session:
        aoi_codes = aoi_code_str.split(", ")
        aoi_names = []
        for aoi_code in aoi_codes:
            aoi_name = session.query(AOI).filter_by(aoi_code=aoi_code).first().aoi_name
            aoi_names.append(aoi_name)
        aoi_name_str = aoi_names[0] if len(aoi_names)==1 else ", ".join(aoi_names)

        return aoi_name_str
    

def get_material_lots(material_names, lots=None):
    with InferSessionLocal() as session:
        conditions = []
        for material_name in material_names:
            material = session.query(Material).filter_by(material_name=material_name).first()
            material_code = material.material_code
            conditions.append(Lot.material_code.like("%" + material_code + "%"))
        if conditions:
            if lots:
                lots = lots.filter(or_(*conditions))
            else:
                lots = session.query(Lot).filter(or_(*conditions))

        return lots
    

def panel_failure_data(panels, session):
    predict_data = []
    for panel in panels:

        lot = session.query(Lot).filter(Lot.lot_code == panel.lot_code).first()

        panel_images = session.query(InferPanelImage).filter(InferPanelImage.panel_code == panel.panel_code).all()

        item = {}
        item["failure_num"] = 0
        for panel_image in panel_images:
            # panel 上去重的缺陷
            detection_info = json.loads(panel_image.detection_info) if panel_image.detection_info else {}

            for key, value in detection_info.items():
                if key not in item:
                    item[key] = 0
                item[key] += len(value)
                item["failure_num"] += len(value)

        item["panel_id"] = panel.panel_id
        item["panel_code"] = panel.panel_code
        item["panel_aoi_time"] = panel.start_time.strftime('%Y-%m-%d %H:%M:%S')
        if lot:
            item["lot_name"] = lot.lot_name
            item["lot_code"] = lot.lot_code
            item["start_time"] = lot.start_time.strftime('%Y-%m-%d %H:%M:%S')
            item["end_time"] = lot.end_time.strftime('%Y-%m-%d %H:%M:%S') if lot.end_time else ""
            material = session.query(Material).filter_by(material_code=panel.material_code).first()
            if material:
                item["material_name"] = material.material_name
            else:
                item["material_name"] = ""
            item["aoi_name"] = session.query(AOI).filter_by(aoi_code=panel.aoi_code).first().aoi_name if panel.aoi_code else ""

        else:
            item["lot_name"] = ""
            item["lot_code"] = ""
            item["start_time"] = ""
            item["end_time"] = ""
            item["material_name"] = ""
        item["ai_result"] = panel.ai_result
        item["aoi_result"] = panel.aoi_result
        item["judgment"] = panel.judgment

        
        # ai 判定结果和人工判定结果做对比, 判断是否复判一致
        # 复判一致
        if item["ai_result"] == item["judgment"]:
            item["judgment_result"] = 1
        # 复判不一致
        else:
            if item["judgment"] == 0:
                item["judgment_result"] = 0
            # 人工复判：NG, 
            elif item["judgment"] == 2:
                # ai ok
                if item["ai_result"] == 1:
                    item["judgment_result"] = 2
                # ai gray
                elif item["ai_result"] == 3:
                    item["judgment_result"] = 1
            # 人工复判 OK
            else:
                item["judgment_result"] = 2

        predict_data.append(item)

    return predict_data


def panel_ok(lot_dir, lot_code, material_code, aoi_code, session):
    for date_time in os.listdir(lot_dir):
        # ================ 开始统计 OK panel 
        # /mnt/md126/jd/lot/192.168.110.4_2025022113051310/20250221
        panel_dir = os.path.join(lot_dir, date_time)
        for panel_name in os.listdir(panel_dir):
            # NG_A2C0375270200250214085021N05AX_0422083858_A2
            if "OK" in panel_name:
                panel_info = panel_name.split("_")

                date_str = panel_info[2]

                current_year = datetime.now().year
                datetime_str = f"{current_year}{date_str[:2]}{date_str[2:4]}{date_str[4:6]}{date_str[6:8]}{date_str[8:10]}"

                panel_time = datetime.strptime(datetime_str, "%Y%m%d%H%M%S")

                # 存储 ok panel
                panel_image = InferPanel(
                    lot_code=lot_code,
                    material_code=material_code,
                    aoi_code=aoi_code,
                    panel_id=panel_name,
                    aoi_result=0,
                    ai_result=0,
                    start_time=panel_time,
                )

                session.add(panel_image)
                session.commit()


def get_infer_type():
    config = Config()
    config.read(os.path.join(root_path, 'config', 'server.ini'))
    infer_type = config.getint('lot', 'infer_type')

    return infer_type


def lot_code_filter(session, **kwargs):
    lot_codes = []

    lots = lot_filter(session, **kwargs)

    if lots:
    
        for lot in lots.all():
            if lot.lot_code not in lot_codes:
                lot_codes.append(lot.lot_code)

    return lot_codes


def panel_filter(session, **kwargs):
    # 以 aoi 生成时间排序
    panel_images = session.query(InferPanel).order_by(InferPanel.start_time.desc())

    if kwargs.get("lot_code"):
        panel_images = panel_images.filter(InferPanel.lot_code==kwargs.get("lot_code"))

    else:
        if kwargs.get("lot_names"):
            lot_codes = lot_code_filter(session, lot_names=kwargs.get("lot_names"))

            panel_images = panel_images.filter(InferPanel.lot_code.in_(lot_codes))

    if kwargs.get("aoi_start_time"):
        aoi_start_time = kwargs.get("aoi_start_time")
        aoi_end_time = kwargs.get("aoi_end_time")
        panel_images = panel_images.filter(and_(InferPanel.start_time >= aoi_start_time, InferPanel.start_time <= aoi_end_time))

    if kwargs.get("panel_id"):
        panel_images = panel_images.filter(InferPanel.panel_id.like("%" + kwargs.get("panel_id") + "%"))

    if kwargs.get("panel_ids"):
        panel_images = panel_images.filter(InferPanel.panel_id.in_(kwargs.get("panel_ids")))

    # 通过 ai 推理的是 ok 还是 ng 过滤
    if kwargs.get("ai_result"):
        panel_images = panel_images.filter(InferPanel.ai_result == kwargs.get("ai_result"))

    if kwargs.get("aoi_result"):
        panel_images = panel_images.filter(InferPanel.aoi_result == kwargs.get("aoi_result"))

    if kwargs.get("start_time"):
        lot_codes = lot_code_filter(session, start_time=kwargs.get("start_time"), end_time=kwargs.get("end_time"))
        panel_images = panel_images.filter(InferPanel.lot_code.in_(lot_codes))

    if kwargs.get("start_time_contain"):
        lot_codes = lot_code_filter(session, start_time_contain=kwargs.get("start_time_contain"), end_time_contain=kwargs.get("end_time_contain"))
        panel_images = panel_images.filter(InferPanel.lot_code.in_(lot_codes))

    if kwargs.get("pjudgment"):
        panel_images = panel_images.filter(InferPanel.judgment == int(kwargs.get("pjudgment")))

    # 过滤人工复判
    if kwargs.get("judgments"):
        conditions = []
        for judgment in kwargs.get("judgments"):
            # 未复判
            if judgment == 0:
                conditions.append(InferPanel.judgment == '0')
            # 复判一致
            elif judgment == 1:
                conditions.append(or_(InferPanel.judgment == InferPanel.ai_result, and_(InferPanel.judgment==2, InferPanel.ai_result==3)))
            # 复判不一致
            elif judgment == 2:
                # bug: 人工 = 2 NG ai_result = 3 gray
                conditions.append(or_(and_(InferPanel.judgment==1, InferPanel.ai_result==2), and_(InferPanel.judgment==2, InferPanel.ai_result==1), and_(InferPanel.judgment==1, InferPanel.ai_result==3)))

        panel_images = panel_images.filter(or_(*conditions))

    # 通过缺陷类型过滤
    if kwargs.get("failure_ids"):
        # 去重
        panel_codes = session.query(InferPanelImage.panel_code).filter(InferPanelImage.failure_id.in_(kwargs.get("failure_ids"))).distinct().all()
        panel_codes = [value for value, in panel_codes]
        panel_images = panel_images.filter(InferPanel.panel_code.in_(panel_codes))

    if kwargs.get("material_names"):
        lot_codes = lot_code_filter(session, material_names=kwargs.get("material_names"))
        panel_images = panel_images.filter(InferPanel.lot_code.in_(lot_codes))

    return panel_images


def lot_filter(session, **kwargs):
    lots = session.query(Lot)

    if kwargs.get("lot_names"):
        lot_names = kwargs.get("lot_names")
        lots = lots.filter(Lot.lot_name.in_(lot_names))

    if kwargs.get("lot_name"):
        lots = lots.filter(Lot.lot_name.like("%" + kwargs.get("lot_name") + "%"))

    # 时间必须包含在里面才可以查询
    if kwargs.get("start_time_contain"):
        start_time = kwargs.get("start_time_contain") 
        end_time = kwargs.get("end_time_contain") 
        lots = lots.filter(Lot.start_time >= start_time).filter(Lot.end_time <= end_time, Lot.end_time.isnot(None))

    # 前端选择的时间和数据库的时间会相差 8 个时区
    # 时间只要有交叉就可以查询
    if kwargs.get("start_time"):
        # 前端虽然发送的是北京时间,但后端接收到的还是 UTC 时间
        start_time = kwargs.get("start_time") 
        end_time = kwargs.get("end_time") 

        time_conditions = []
        time_conditions.append(and_(Lot.start_time <= end_time, Lot.start_time >= start_time))
        time_conditions.append(and_(Lot.end_time <= end_time, Lot.end_time >= start_time, Lot.end_time.isnot(None)))
        time_conditions.append(and_(Lot.end_time >= end_time, Lot.start_time <= start_time, Lot.end_time.isnot(None)))

        lots = lots.filter(or_(*time_conditions))

    if kwargs.get("material_names"):
        conditions = []
        for material_name in kwargs.get("material_names"):
            material = session.query(Material).filter_by(material_name=material_name).first()
            if material:
                material_code = material.material_code
                conditions.append(Lot.material_code.like("%" + material_code + "%"))
        if conditions:
            lots = lots.filter(or_(*conditions))
        else:
            lots = lots.filter(Lot.material_code.in_([]))

    if kwargs.get("material_name"):
        material = session.query(Material).filter_by(material_name=material_name).first()
        if material:
            material_code = material.material_code
            lots = lots.filter(Lot.material_code.like("%" + material_code + "%"))

        else:
            lots = None
    
    return lots


def failure_filter(session, **kwargs):
    failure_images = session.query(InferPanelImage)

    if kwargs.get("panel_code") and kwargs.get("lot_code"):
        failure_images = failure_images.filter(InferPanelImage.panel_code==kwargs.get("panel_code"))

    else:
        if kwargs.get("lot_code"):
            failure_images = failure_images.filter(InferPanelImage.lot_code==kwargs.get("lot_code"))
        
        if kwargs.get("lot_name"):
            lot_codes = lot_code_filter(session, lot_name=kwargs.get("lot_name"))
            failure_images = failure_images.filter(InferPanelImage.lot_code.in_(lot_codes))

        if kwargs.get("panel_id"):
            panel_codes = panel_code_filter(session, panel_id=kwargs.get("panel_id"))
            failure_images = failure_images.filter(InferPanelImage.panel_code.in_(panel_codes))

    if kwargs.get("start_time"):
        lot_codes = lot_code_filter(session, start_time=kwargs.get("start_time"), end_time=kwargs.get("end_time"))
        failure_images = failure_images.filter(InferPanelImage.lot_code.in_(lot_codes))

    if kwargs.get("failure_ids"):
        failure_images = failure_images.filter(InferPanelImage.failure_id.in_(kwargs.get("failure_ids")))

    if kwargs.get("judgments"):
        panel_codes = panel_code_filter(session, judgments=kwargs.get("judgments"))
        failure_images = failure_images.filter(InferPanelImage.panel_code.in_(panel_codes))

    return failure_images
    

def panel_code_filter(session, **kwargs):

    panel_codes = []

    panel_images = panel_filter(session, **kwargs)

    for panel_image in panel_images.all():
        panel_codes.append(panel_image.panel_code)

    return panel_codes


def get_material_parameter(ini_path, section_name, conf_name):
    config = Config()
    config.read(os.path.join(ini_path))
    res = config.get(section_name, conf_name)

    return res



